#include "algoritm.h"
#include "complex.h"
// #include "define.h"
// #include "SimpleIni.h"
#include <math.h>
#include <stdio.h>
#include "fft_s.h"
double Ienergy(complex *A, int N)
{
    double energy = 0.;
    for (int i = 0; i < N; i++)
    {
        energy += norm(A[i]);
    }
    return energy;
}
double Ipuls(complex *A, int N)
{
    double puls;
    for (int i = 1; i < N-1; i++)
    {
        puls += imag(A[i]*conj(A[i+1] - A[i-1]));
    }
    return puls;
}
double Ihamilt(complex *A, int N, double hx, double Y, double D)
{
    double normAi, hamilt = 0, hamilt_tmp = 0.;
    for (int i = 0; i < N; i++)
    {
        normAi = norm(A[i]);
        hamilt += normAi*normAi;
    }
    for (int i = 0; i < N-1; i++)
    {
        hamilt_tmp += norm(A[i+1] - A[i]);
    }
    return hamilt * Y/2 - hamilt_tmp * D/hx;
}

void fft_algoritm(parameters par)
{
    complex *A = new complex[par.Nx];
    complex *Aprev = new complex[par.Nx];
    complex *Anew = new complex[par.Nx];
    complex *swap;
    complex *Fprev = new complex[par.Nx];
    complex *Fnew = new complex[par.Nx];
    complex *Fnlprev = new complex[par.Nx];
    complex *Fnlnew = new complex[par.Nx];
    complex *AnormA2 = new complex[par.Nx];

    double *lambda_2 = new double[par.Nx];
    fftdata *Afft = new fftdata(par.Nx);

    double hx = par.L/par.Nx, hz = par.Lz/par.Nz;
    complex iDhz_2 = complex(0,1) * par.D * hz / 2;
    complex iYhz_2 = complex(0,1) * par.Y * hz / 2;

    for (int i=0; i<par.Nx; i++)
    {
        int xi = i - par.X0I;
        double x2 = (double) (xi * xi) * (hx * hx);
        A[i]=exp( - x2 / (par.R*par.R));
        // printf("%f\n", A[i].real());
    }
    
    for (int i = 0; i < par.Nx; i++)
        Aprev[i] = A[i];

    fftr_rep(A, Fprev, Afft);

    for (int i = 0; i < par.Nx/2; i++)
    {
        lambda_2[i] = 2 * M_PI * i/par.L;
        lambda_2[i] = lambda_2[i] * lambda_2[i];
    }
    for (int i = par.Nx/2; i < par.Nx; i++)
    {
        lambda_2[i] = 2 * M_PI * (i-par.Nx)/par.L;
        lambda_2[i] = lambda_2[i] * lambda_2[i];
    }
    
    double energy0, puls0, hamilt0;
    energy0 = Ienergy(A, par.Nx);
    puls0 = Ipuls(A, par.Nx);
    hamilt0 = Ihamilt(A, par.Nx, hx, par.Y, par.D);
    // printf("%e\n", energy0);
    for (int j = 0; j < par.Nz; j++)
    {
        int itteraton = 0;
        bool flag = true; 
        for (int i = 0; i < par.Nx; i++)
        {
            AnormA2[i] = norm(Aprev[i]) * Aprev[i];
        }
        
        fftr_rep(Aprev, Fprev, Afft);
        fftr_rep(AnormA2, Fnlprev, Afft);
        for (int i = 0; i < par.Nx; i++)
        {
            // Fnew[i] = Fprev[i];
            Fnlnew[i] = Fnlprev[i]; 
        }
        
        while (flag)
        {
            // printf("\ritteraton %i\n", itteraton);
            for(int k=1; k<par.Nx; k++)
            {
                Fnew[k] = (Fprev[k] * (1 - iDhz_2 * lambda_2[k]) + iYhz_2 * (Fnlprev[k] + Fnlnew[k])) / (1 + iDhz_2 * lambda_2[k]);
            }
            fft_rep(Fnew, Anew, Afft);

            flag = false;
            for (int i = par.Nx-1; i > 0; i--)
            {
                double diff = abs(Anew[i] - Aprev[i]);
                if (diff > par.EPSz * abs(Aprev[i]))
                {
                    flag = true;
                }
            }

            for (int i = 0; i < par.Nx; i++)
            {
                AnormA2[i] = norm(Anew[i]) * Anew[i];
            }
            fftr_rep(AnormA2, Fnlnew, Afft);


            // printf("\r%2i real %f im %f", itteraton,  Anew[par.Nx/2].real(), Anew[par.Nx/2].imag());
            // printf("%i", flag);
            swap = Anew;
            Anew = Aprev;
            Aprev = swap;
            
            // if(flag) break;
            itteraton++;
            // printf("\ritteratons %i\n", itteraton);
        }
        // printf("\r%2i", itteraton);
        for (int i = 0; i < par.Nx; i++)
            A[i] = Aprev[i];
        // fftr_rep(Aprev, Fprev, Afft);
        printf("\r progress %i", j);

    }
    // printf("%e\n", energy0);
    double energy = 0., puls = 0., hamilt = 0.;
    energy = Ienergy(A, par.Nx);
    puls = Ipuls(A, par.Nx);
    hamilt = Ihamilt(A, par.Nx, hx, par.Y, par.D);
    printf("\n");
    printf("%e %e\n", energy, (energy-energy0)/energy0);
    printf("%e %e\n", puls, puls0);
    printf("%e %e\n", hamilt, (hamilt-hamilt0)/hamilt0);
    delete(A);
    delete(Anew);
    delete(Aprev);
    delete(Fprev);
    delete(Fnew);
    delete(Fnlprev);
    delete(Fnlnew);
    delete(AnormA2);
    delete(lambda_2);
    delete(Afft);
}

void sweep_algoritm(parameters par)
{
    complex *A = new complex[par.Nx];
    complex *ASprev = new complex[par.Nx];
    complex *ASnew = new complex[par.Nx];
    complex *Aswap;
    complex *alpha = new complex[par.N];
    complex *beta = new complex[par.N];

    double hx = par.L/par.N, hz = par.Lz/par.Nz;
    complex i2hx2_Dhz2 = (2*complex(0,1)*hx*hx/(par.D*hz*hz));
    double Yhx2_D = par.Y * hx*hx / par.D;

    for(int i=0; i<par.Nx; i++)
    {
        int xi = i - par.X0I;
        double x2 = (double) (xi * xi) * (hx * hx);
        A[i]=exp( - x2 / (par.R*par.R));
        // printf("%f\n", A[i].real());
    }
    alpha[0] = 0;
    beta[0] = 0;
    for(int i=1; i<par.N; i++)
    {
        alpha[i] = 1. / (2. - i2hx2_Dhz2 - alpha[i-1]);
    }

    for (int i = 0; i < par.Nx; i++)
        ASprev[i] = A[i];
    
    double energy0, puls0, hamilt0;
    energy0 = Ienergy(A, par.Nx);
    puls0 = Ipuls(A, par.Nx);
    hamilt0 = Ihamilt(A, par.Nx, hx, par.Y, par.D);
    // printf("%e\n", energy0);
    for (int j = 0; j < par.Nz; j++)
    {
        // itteration while
        int itteraton = 0;
        bool flag = true; 
        while (flag)
        {
            // printf("\ritteraton %i\n", itteraton);
            for(int i=1; i<par.N; i++)
            {
                // alpha[i] = 1. / (2. - i2hx2_Dhz2 - alpha[i-1]);
                beta[i] = alpha[i] * (beta[i-1] + A[i+1] - (2. + i2hx2_Dhz2) * A[i] + A[i-1] + Yhx2_D * (ASprev[i] * norm(ASprev[i]) + A[i] * norm(A[i])));
            }
            flag = false;
            for (int i = par.Nx-1; i > 0; i--)
            {
                ASnew[i] = ASnew[i+1] * alpha[i] + beta[i];
                double diff = abs(ASnew[i] - ASprev[i]);
                if (diff > par.EPSz * abs(ASprev[i]))
                {
                    flag = true;
                }
            }

            // printf("%2i real %f im %f\n", itteraton,  ASnew[Nx/2].real(), ASnew[Nx/2].imag());
            // printf("%i", flag);
            Aswap = ASnew;
            ASnew = ASprev;
            ASprev = Aswap;
            
            // if(flag) break;
            itteraton++;
            // printf("\ritteratons %i\n", itteraton);
        }
        for (int i = 0; i < par.Nx; i++)
            A[i] = ASprev[i];
    }
    // printf("%e\n", energy0);
    double energy, puls, hamilt;
    energy = Ienergy(A, par.Nx);
    puls = Ipuls(A, par.Nx);
    hamilt = Ihamilt(A, par.Nx, hx, par.Y, par.D);

    printf("%e %e\n", energy, (energy-energy0)/energy0);
    printf("%e %e\n", puls, puls0);
    printf("%e %e\n", hamilt, (hamilt-hamilt0)/hamilt0);
    delete(A);
    delete(ASnew);
    delete(ASprev);
    delete(alpha);
    delete(beta);
}
